from PIL import Image
import numpy as np
from tensorflow.keras.models import load_model
import cv2

# 加载模型
model = load_model('./model/lenet5.h5')

def predict(image_path):
    # 以黑白方式读取图片
    img = Image.open(image_path).convert('L')
    img.resize((28,28))
    img = np.reshape(img, (28, 28, 1)) / 255.
    x = np.array([1 - img])

    y = model.predict(x)

    # 因为x只传入了一张图片，取y[0]即可
    # np.argmax()取得最大值的下标，即代表的数字
    print(image_path)
    print(y[0])
    print('the number is :', np.argmax(y[0]))


if __name__ == "__main__":
    img = cv2.imread('./imgs/3.jpg', 0)
    img = cv2.resize(img, (28, 28))
    cv2.imwrite('./test.jpg', img)
    predict('./test.jpg')